from flask import Flask, request, Response
from langchain_community.llms import Ollama
import os
from flask_cors import CORS  # 导入 CORS
import json
from langchain_ollama import ChatOllama

app = Flask(__name__)
CORS(app)  # 启用 CORS

model = "qwen2.5:7b"

# 初始化 Ollama 模型，设置格式化参数
# llm = Ollama(
#     model=model, 
#     base_url="http://127.0.0.1:11434",
#     format="json"  # 尝试使用 JSON 格式
# )
llm = ChatOllama(model="qwen2.5:7b", temperature=0.7)

# 定义流式响应函数
def generate_response(question):
    try:
        # 方案1：使用 JSON 格式化
        # response = llm.invoke(
        #     question, 
        #     stream=True,
        #     stop=None,  # 不使用停止标记
        #     temperature=0.7
        #     # response_format={"type": "text"}  # 指定响应格式
        # )
        
        for chunk in llm.stream(question):
            try:
                # 尝试解析为 JSON
                if isinstance(chunk, dict):
                    content = chunk.get('content', chunk)
                else:
                    content = str(chunk)
                
                # 确保换行符被保留
                content = content.replace('\\n', '\n')
                content = content.replace('\\r', '\n')

                arr = content.split(" ")
                content = arr[0].split("=")[1]
                content = content.replace('\'', "")
                
                print(content)
                # print(f"{content}")
                
                if(content) :
                    # 使用 JSON 编码确保特殊字符被正确处理
                    formatted_content = json.dumps({"model":f"{model}","created_at":"","message":{"role":"assistant","content":f"{content}"},"done":False})
                    yield f"{formatted_content}\n"                 

            except Exception as e:
                print(f"1-异常：{str(e)}")
                # 如果 JSON 解析失败，直接返回原始内容
                yield f"data: {json.dumps({'content': str(chunk)})}\n\n"
            
    except Exception as e:
        print(f"2-异常：{str(e)}")
        yield f"data: {json.dumps({'error': str(e)})}\n\n"

@app.route('/ask', methods=['POST'])
def ask():
    data = request.get_json()
    question = data.get('messages', '')
    
    if not question:
        return {"status": "error", "message": "问题不能为空"}, 400
    
    return Response(generate_response(question), mimetype='text/event-stream')

if __name__ == "__main__":
    app.run(host='127.0.0.1', port=8001) 